import numpy as np
from scipy.integrate import quad

# 更正后的 Hermite 基函数二阶导数定义


def ddphi_0(x):
    return -6 + 12 * x  # phi_0 的二阶导数


def ddphi_1(x):
    return -4 + 6 * x   # phi_1 的二阶导数


def ddphi_2(x):
    return 6 - 12 * x   # phi_2 的二阶导数


def ddphi_3(x):
    return -2 + 6 * x   # phi_3 的二阶导数


class Mesh:
    def __init__(self, n: int, interval: list = [0, 1]):
        self.mesh = np.linspace(interval[0], interval[1], n+1)
        self.h = (interval[1] - interval[0]) / n
        self.n = n
        self.n_nodes = n + 1


class FE:
    def __init__(self, mesh: Mesh):
        self.mesh = mesh.mesh
        self.h = mesh.h
        self.n = mesh.n
        self.n_nodes = mesh.n_nodes
        self.n_dofs_per_node = 2
        self.n_dofs = self.n_dofs_per_node * self.n_nodes
        self.K = np.zeros((self.n_dofs, self.n_dofs))
        self.b = np.zeros(self.n_dofs)

    def assemble_matrix(self):
        n_dofs_per_element = 4  # 每个单元4个自由度
        for i in range(self.n):
            Ke = np.zeros((n_dofs_per_element, n_dofs_per_element))
            for alpha in range(n_dofs_per_element):
                # 选择二阶导数函数
                if alpha == 0:
                    dd_alpha = ddphi_0
                elif alpha == 1:
                    dd_alpha = ddphi_1
                elif alpha == 2:
                    dd_alpha = ddphi_2
                else:
                    dd_alpha = ddphi_3

                for beta in range(n_dofs_per_element):
                    if beta == 0:
                        dd_beta = ddphi_0
                    elif beta == 1:
                        dd_beta = ddphi_1
                    elif beta == 2:
                        dd_beta = ddphi_2
                    else:
                        dd_beta = ddphi_3

                    # 积分项调整为二阶导数乘积及正确缩放因子
                    def integrand(x): return (1 / self.h**3) * \
                        dd_alpha(x) * dd_beta(x)
                    Ke[alpha, beta] = quad(integrand, 0, 1)[0]

            # 组装到全局矩阵
            dof_indices = [2*i, 2*i+1, 2*(i+1), 2*(i+1)+1]
            for a in range(4):
                for b in range(4):
                    self.K[dof_indices[a], dof_indices[b]] += Ke[a, b]

    # 保留原有assemble_vector和边界条件处理...

    def dirichlet_bc(self, left_value, right_value):
        # 左端u和右端u的值固定
        left_dof = 0
        right_dof = 2 * (self.n_nodes - 1)
        self.K[left_dof, :] = 0
        self.K[left_dof, left_dof] = 1
        self.K[right_dof, :] = 0
        self.K[right_dof, right_dof] = 1
        self.b[left_dof] = left_value
        self.b[right_dof] = right_value

    def neumann_bc(self, left_deriv=0, right_deriv=0):
        # 左端u'和右端u'固定
        left_deriv_dof = 1
        right_deriv_dof = 2 * (self.n_nodes - 1) + 1
        self.K[left_deriv_dof, :] = 0
        self.K[left_deriv_dof, left_deriv_dof] = 1
        self.K[right_deriv_dof, :] = 0
        self.K[right_deriv_dof, right_deriv_dof] = 1
        self.b[left_deriv_dof] = left_deriv
        self.b[right_deriv_dof] = right_deriv

    def solve(self, f, left_value, right_value, left_deriv=0, right_deriv=0):
        self.assemble_matrix()
        self.assemble_vector(f)
        self.dirichlet_bc(left_value, right_value)
        self.neumann_bc(left_deriv, right_deriv)
        self.u = np.linalg.solve(self.K, self.b)
        return self.u
